# MNE-Python 和 YASA 结合的睡眠分析程序
# 功能：前30秒提供实时指标，之后每3秒提供一次滑动窗口分析结果
# 优势：平衡了实时性、科学准确性与用户体验

import numpy as np
import mne
import yasa
import time
import csv
import pickle
import os
import sys
from datetime import datetime

# 添加父目录到路径以导入NeuroPy3
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from NeuroPy3 import NeuroPy 

class SleepAnalyzer:
    def __init__(self, port="COM5", sfreq=512, log_step_sec=1, analysis_step_sec=3, window_sec=300):
        self.port = port
        self.sfreq = sfreq
        self.log_step_size = sfreq * log_step_sec  # 每秒记录日志
        self.analysis_step_size = sfreq * analysis_step_sec  # 每3秒分析一次
        self.window_size = sfreq * window_sec  # 5分钟窗口
        
        self.log_counter = 0  # 日志计数器
        self.analysis_counter = 0  # 分析计数器
        
        self.buffer = []
        self.poor_signal_buffer = []
        
        self.neuropy = NeuroPy(port, 57600)
        self.log_file = "sleep_analysis_log.csv"
        
        # 5分钟数据存储文件路径
        self.data_cache_file = "5min_data_cache.pkl"
        
        self.info = mne.create_info(ch_names=['NeuroSky'], sfreq=self.sfreq, ch_types=['eeg'])
        
        # 平滑处理相关属性
        self.stage_history = []  # 存储最近的睡眠阶段预测结果
        self.smooth_window = 3   # 平滑窗口大小（最近3次预测）
        
        self._init_log()
        
    def _smooth_stage_prediction(self, raw_stage):
        """平滑处理睡眠阶段预测结果
        使用滑动窗口多数投票机制减少不合理的阶段跳跃
        """
        # 将新预测添加到历史记录
        self.stage_history.append(raw_stage)
        
        # 保持历史记录在窗口大小内
        if len(self.stage_history) > self.smooth_window:
            self.stage_history = self.stage_history[-self.smooth_window:]
        
        # 如果历史记录不足，直接返回原始预测
        if len(self.stage_history) < 2:
            return raw_stage
        
        # 统计各阶段出现次数
        stage_counts = {}
        for stage in self.stage_history:
            stage_counts[stage] = stage_counts.get(stage, 0) + 1
        
        # 返回出现次数最多的阶段（多数投票）
        smoothed_stage = max(stage_counts, key=stage_counts.get)
        
        # 如果平滑后的结果与原始预测不同，输出提示
        if smoothed_stage != raw_stage:
            print(f"平滑处理: {raw_stage} -> {smoothed_stage} (基于最近{len(self.stage_history)}次预测)")
        
        return smoothed_stage
        
    def _init_log(self):
        """初始化CSV日志文件"""
        headers = ['时间戳', '数据长度(秒)', '睡眠阶段', 'Delta功率(dB)', 'Theta功率(dB)', 'Alpha功率(dB)', 'Beta功率(dB)', 'Gamma功率(dB)', '信号质量']
        with open(self.log_file, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(headers)
    
    def _data_callback(self, raw_value):
        """接收原始EEG数据的回调函数"""
        self.buffer.append(raw_value)
    
    def _poor_signal_callback(self, poor_signal_value):
        """信号质量回调函数"""
        self.poor_signal_buffer.append(poor_signal_value)
    
    def _process_data(self):
        """处理缓冲区数据并进行分析"""
        buffer_length = len(self.buffer)
        
        # --- 1. 每秒记录日志 ---
        if buffer_length >= self.log_step_size * (self.log_counter + 1):
            self.log_counter += 1
            
            # 提取最新1秒数据进行实时指标计算
            start_idx = (self.log_counter - 1) * self.log_step_size
            end_idx = self.log_counter * self.log_step_size
            log_data = np.array(self.buffer[start_idx:end_idx])
            
            log_data_volts = log_data / 1e6
            raw_log = mne.io.RawArray([log_data_volts], self.info, verbose=False)
            
            # 计算实时频谱指标
            psds, freqs = raw_log.compute_psd(method='welch', fmin=1.0, fmax=40.0, verbose=False).get_data(return_freqs=True, picks='NeuroSky')
            psds_db = 10 * np.log10(psds)
            
            delta_power = self._get_band_power(psds_db, freqs, 1, 4)
            theta_power = self._get_band_power(psds_db, freqs, 4, 8)
            alpha_power = self._get_band_power(psds_db, freqs, 8, 13)
            beta_power = self._get_band_power(psds_db, freqs, 13, 30)
            gamma_power = self._get_band_power(psds_db, freqs, 30, 40)
            
            signal_quality = np.mean(self.poor_signal_buffer) if self.poor_signal_buffer else 0
            
            # 默认睡眠阶段
            stage = "数据累积中" if buffer_length < self.window_size else "等待分析"
            
            # 记录每秒日志
            print(f"第{self.log_counter}秒 - 缓冲区: {buffer_length/self.sfreq:.1f}秒, Delta: {delta_power:.2f} dB")
            self._log_result(1.0, stage, delta_power, theta_power, alpha_power, beta_power, gamma_power, signal_quality)
        
        # --- 2. 每3秒进行YASA分析（仅在5分钟后） ---
        if buffer_length >= self.window_size and buffer_length >= self.analysis_step_size * (self.analysis_counter + 1):
            self.analysis_counter += 1
            
            # 保存5分钟数据到缓存文件
            self._save_5min_data_cache()
            
            # 进行YASA睡眠分期
            stage = self._perform_yasa_analysis()
            
            print(f"=== 分析完成 ===\n")
        
        # --- 3. 滑动缓冲区管理 ---
        # 保持缓冲区在合理大小，删除过旧的数据
        if buffer_length > self.window_size * 1.2:  # 超过5分钟的1.2倍时清理
            excess = buffer_length - self.window_size
            del self.buffer[:excess]
            # 相应调整计数器
            removed_seconds = excess // self.sfreq
            self.log_counter = max(0, self.log_counter - removed_seconds)
            self.analysis_counter = max(0, self.analysis_counter - removed_seconds // 3)
            print(f"清理缓冲区，删除{removed_seconds}秒旧数据")
        
        # 清理信号质量缓冲区
        if len(self.poor_signal_buffer) > 10:
            self.poor_signal_buffer = self.poor_signal_buffer[-10:]

    def _get_band_power(self, psds, freqs, low_freq, high_freq):
        """从完整的PSD数组中提取指定频段的平均功率（dB单位）"""
        band_idx = np.logical_and(freqs >= low_freq, freqs <= high_freq)
        return psds[0, band_idx].mean()

    def _save_5min_data_cache(self):
        """保存当前5分钟数据到缓存文件"""
        try:
            # 使用最新的5分钟数据
            window_data = self.buffer[-self.window_size:] if len(self.buffer) >= self.window_size else self.buffer
            cache_data = {
                'buffer': window_data.copy(),
                'timestamp': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'sfreq': self.sfreq
            }
            with open(self.data_cache_file, 'wb') as f:
                pickle.dump(cache_data, f)
            print(f"已保存5分钟数据到缓存文件: {self.data_cache_file}")
        except Exception as e:
            print(f"保存数据缓存失败: {e}")
    
    def _load_5min_data_cache(self):
        """从缓存文件加载5分钟数据"""
        try:
            if os.path.exists(self.data_cache_file):
                with open(self.data_cache_file, 'rb') as f:
                    cache_data = pickle.load(f)
                print(f"已从缓存加载5分钟数据 (时间戳: {cache_data['timestamp']})")
                return cache_data
            else:
                print("缓存文件不存在")
                return None
        except Exception as e:
            print(f"加载数据缓存失败: {e}")
            return None
    
    def _perform_yasa_analysis(self, use_cached_data=False):
        """执行YASA睡眠分期分析"""
        print(f"\n=== 第{self.analysis_counter}次YASA分析 ===")
        
        # 选择数据源
        if use_cached_data:
            cached_data = self._load_5min_data_cache()
            if cached_data is None:
                print("无法加载缓存数据，使用当前缓冲区数据")
                data_buffer = self.buffer[-self.window_size:] if len(self.buffer) >= self.window_size else self.buffer
            else:
                data_buffer = cached_data['buffer']
        else:
            data_buffer = self.buffer[-self.window_size:] if len(self.buffer) >= self.window_size else self.buffer
        
        print(f"缓冲区长度: {len(data_buffer)} 样本 ({len(data_buffer)/self.sfreq:.1f}秒)")
        
        try:
            # 使用数据进行YASA分析
            window_data = np.array(data_buffer)
            window_data_volts = window_data / 1e6
            window_raw = mne.io.RawArray([window_data_volts], self.info, verbose=False)
            
            filtered_raw = window_raw.copy().filter(l_freq=0.3, h_freq=35, verbose=False)
            
            # 使用YASA的SleepStaging类进行自动睡眠分期
            sls = yasa.SleepStaging(filtered_raw, eeg_name='NeuroSky')
            hypno_result = sls.predict()
            
            # 获取睡眠阶段数组 - hypno_result本身就是一个对象，有hypno属性
            if hasattr(hypno_result, 'hypno'):
                hypno = hypno_result.hypno
            else:
                # 如果hypno_result直接是数组
                hypno = hypno_result
            
            # 获取最后一个epoch的睡眠阶段
            stage_mapping = {0: 'Wake', 1: 'N1', 2: 'N2', 3: 'N3', 4: 'REM', 'W': 'Wake', 'N1': 'N1', 'N2': 'N2', 'N3': 'N3', 'R': 'REM'}
            if len(hypno) > 0:
                stage_num = hypno.iloc[-1] if hasattr(hypno, 'iloc') else hypno[-1]
                raw_stage = stage_mapping.get(stage_num, f'Unknown({stage_num})')
                print(f"YASA 原始预测: {raw_stage} (数值: {stage_num})")
                
                # 应用平滑处理
                smoothed_stage = self._smooth_stage_prediction(raw_stage)
                print(f"最终睡眠阶段: {smoothed_stage}")
                
                # 更新最近的日志记录中的睡眠阶段（使用平滑后的结果）
                self._update_recent_logs_with_stage(smoothed_stage)
                return smoothed_stage
            else:
                print("YASA 睡眠分期结果为空")
                return "未知"
            
        except Exception as e:
            print(f"YASA 睡眠分期失败: {e}")
            return "分析失败"
    
    def test_yasa_with_cached_data(self):
        """使用缓存的5分钟数据测试YASA分析"""
        print("\n=== 使用缓存数据测试YASA分析 ===")
        stage = self._perform_yasa_analysis(use_cached_data=True)
        print(f"测试结果: {stage}")
        return stage
    
    def _update_recent_logs_with_stage(self, stage):
        """更新最近的日志记录中的睡眠阶段"""
        # 读取现有日志文件并更新最近3秒的记录
        try:
            import pandas as pd
            df = pd.read_csv(self.log_file, encoding='utf-8')
            if len(df) >= 3:
                # 更新最近3行的睡眠阶段
                df.iloc[-3:, df.columns.get_loc('睡眠阶段')] = stage
                df.to_csv(self.log_file, index=False, encoding='utf-8')
                print(f"已更新最近3秒日志的睡眠阶段为: {stage}")
        except Exception as e:
            print(f"更新日志失败: {e}")
            print(f"当前睡眠阶段: {stage}")
    
    def _log_result(self, length, stage, delta, theta, alpha, beta, gamma, quality):
        """记录分析结果到日志"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        with open(self.log_file, 'a', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow([timestamp, length, stage, f"{delta:.6f}", f"{theta:.6f}", f"{alpha:.6f}", f"{beta:.6f}", f"{gamma:.6f}", f"{quality:.2f}"])
    
    def start(self, preload_cache=True):
        """启动数据采集
        
        Args:
            preload_cache: 是否预加载缓存数据作为初始缓冲区
        """
        # 如果存在缓存且选择预加载，则加载到缓冲区
        if preload_cache and os.path.exists(self.data_cache_file):
            try:
                cached_data = self._load_5min_data_cache()
                if cached_data and 'buffer' in cached_data:
                    self.buffer = cached_data['buffer'].copy()
                    print(f"已预加载缓存数据到缓冲区，长度: {len(self.buffer)} 样本 ({len(self.buffer)/self.sfreq:.1f}秒)")
                    print("现在可以立即开始YASA分析！")
            except Exception as e:
                print(f"预加载缓存数据失败: {e}，将从空缓冲区开始")
        
        self.neuropy.setCallBack("rawValue", self._data_callback)
        self.neuropy.setCallBack("poorSignal", self._poor_signal_callback)
        self.neuropy.start()
        print(f"设备已连接到 {self.port}，开始采集数据...")
        
        # 等待缓冲区累积到至少一个步长的数据
        while len(self.buffer) < self.log_step_size:
            time.sleep(0.1)
        
        try:
            while True:
                self._process_data()
                time.sleep(0.1)  # 短暂休眠避免CPU占用过高
        except KeyboardInterrupt:
            print("\n停止采集...")
        finally:
            self.neuropy.stop()
            print("设备已断开")

if __name__ == "__main__":
    # 每秒记录日志，每3秒分析一次，使用5分钟窗口
    analyzer = SleepAnalyzer(log_step_sec=1, analysis_step_sec=1, window_sec=300)
    
    # 提供启动选项
    print("=== EEG睡眠分析器启动选项 ===")
    
    if os.path.exists(analyzer.data_cache_file):
        print("检测到缓存的5分钟数据文件")
        print("请选择启动模式：")
        print("1. 使用缓存数据开始并接受新数据 (推荐)")
        print("2. 不使用缓存，从头开始录取数据")
        print("3. 仅测试缓存数据 (不录取新数据)")
        
        while True:
            user_input = input("请输入选项 (1/2/3): ").strip()
            
            if user_input == '1':
                print("\n选择：使用缓存数据开始并接受新数据")
                print("将预加载缓存数据到缓冲区，然后开始实时采集...")
                analyzer.start(preload_cache=True)
                break
            elif user_input == '2':
                print("\n选择：不使用缓存，从头开始录取数据")
                print("将清空缓冲区，从零开始采集数据...")
                analyzer.start(preload_cache=False)
                break
            elif user_input == '3':
                print("\n选择：仅测试缓存数据")
                print("正在使用缓存数据进行YASA分析...")
                analyzer.test_yasa_with_cached_data()
                break
            else:
                print("无效输入，请输入 1、2 或 3")
    else:
        print("未检测到缓存数据文件")
        print("将从空缓冲区开始采集数据...")
        analyzer.start(preload_cache=False)